package server

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"net"
	"net/http"
	"port-forwarding/config"
	"strconv"
	"time"
)

type JsonResult struct {
	Code int    `json:"code"`
	Msg  string `json:"msg"`
}

func ListenAndServe(listenAddr string, forwardAddr string, config *config.Config) {
	ln, err := net.Listen("tcp", listenAddr)
	if err != nil {
		log.Fatalf("listen addr %s failed: %s", listenAddr, err.Error())
	}

	d := net.Dialer{Timeout: time.Second * 10}

	_, err = d.Dial("tcp", forwardAddr)
	if err != nil {
		log.Printf("try connect forward address %s failed: %s\n", forwardAddr, err.Error())
		return
	}

	log.Printf("accept %s to %s\n", listenAddr, forwardAddr)

	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Printf("accept %s error: %s\n", listenAddr, err.Error())
		}

		// check white url

		isWhite := config.CheckWhiteUrl(conn.RemoteAddr().String())
		if !isWhite {
			// 使用自定义的 ResponseWriter
			crw := NewCustomResponseWriter(conn)
			fmt.Printf("connection request addr is not in white list: %s\n", conn.RemoteAddr().String())
			marshal, _ := json.Marshal(JsonResult{Code: 400, Msg: fmt.Sprintf("request not allow")})
			// 设置 HTTP 状态码
			crw.WriteHeader(http.StatusForbidden)
			// 写入 JSON 数据
			crw.Write(marshal)
			//conn.Write(marshal)
			// 冲刷缓冲区，将数据发送到连接
			err = crw.Flush()
			if err != nil {
				panic(err)
			}
			continue
		}

		go HandleRequest(conn, forwardAddr)
	}
}

// CustomResponseWriter 实现了 http.ResponseWriter 接口
type CustomResponseWriter struct {
	conn         net.Conn
	statusCode   int
	headerBuffer bytes.Buffer
	bodyBuffer   bytes.Buffer
}

// NewCustomResponseWriter 创建一个新的 CustomResponseWriter
func NewCustomResponseWriter(conn net.Conn) *CustomResponseWriter {
	return &CustomResponseWriter{
		conn:         conn,
		statusCode:   http.StatusOK,
		headerBuffer: bytes.Buffer{},
		bodyBuffer:   bytes.Buffer{},
	}
}

// Header 返回当前的头部信息
func (crw *CustomResponseWriter) Header() http.Header {
	return http.Header{}
}

// WriteHeader 设置 HTTP 响应的状态码
func (crw *CustomResponseWriter) WriteHeader(code int) {
	crw.statusCode = code
}

// Write 写入响应体
func (crw *CustomResponseWriter) Write(b []byte) (int, error) {
	return crw.bodyBuffer.Write(b)
}

// Flush 将数据写入 net.Conn
func (crw *CustomResponseWriter) Flush() error {
	// 构建 HTTP 响应头部
	crw.headerBuffer.WriteString("HTTP/1.1 ")
	crw.headerBuffer.WriteString(strconv.Itoa(crw.statusCode))
	crw.headerBuffer.WriteString(" OK\r\n")

	// 添加 Content-Type
	crw.headerBuffer.WriteString("Content-Type: application/json\r\n")

	// 添加 Content-Length
	crw.headerBuffer.WriteString("Content-Length: ")
	crw.headerBuffer.WriteString(strconv.Itoa(len(crw.bodyBuffer.Bytes())))
	crw.headerBuffer.WriteString("\r\n")

	// 添加空行
	crw.headerBuffer.WriteString("\r\n")

	// 写入头部和体
	_, err := crw.conn.Write(crw.headerBuffer.Bytes())
	if err != nil {
		return err
	}

	_, err = crw.conn.Write(crw.bodyBuffer.Bytes())
	return err
}
